#ifndef DDGI_COMMON_GLSL
#define DDGI_COMMON_GLSL

#ifndef M_PI
#define M_PI 3.14159265359
#endif

struct DDGIPayload
{
    vec3 L;
    vec3 T;
    float hitDistance;
    RNG rng;
};

struct DDGIUniforms
{
    vec3 gridStartPosition;
    vec3 gridStep;
    ivec3 probeCounts;
    float maxDistance;
    float depthSharpness;
    float hysteresis;
    float normalBias;
    float energyPreservation;
    int irradianceProbeSideLength;
    int irradianceTextureWidth;
    int irradianceTextureHeight;
    int depthProbeSideLength;
    int depthTextureWidth;
    int depthTextureHeight;
    int raysPerProbe;
    int visibilityTest;
};



ivec3 BaseGridCoord(in DDGIUniforms ddgi, vec3 X) 
{
    return clamp(ivec3((X - ddgi.gridStartPosition) / ddgi.gridStep), ivec3(0, 0, 0), ivec3(ddgi.probeCounts) - ivec3(1, 1, 1));
}



vec3 GridCoordToPosition(in DDGIUniforms ddgi, ivec3 c)
{
    return ddgi.gridStep * vec3(c) + ddgi.gridStartPosition;
}



int GridCoordToProbeIndex(in DDGIUniforms ddgi, in ivec3 probeCoords) 
{
    return int(probeCoords.x + probeCoords.y * ddgi.probeCounts.x + probeCoords.z * ddgi.probeCounts.x * ddgi.probeCounts.y);
}



ivec3 ProbeIndexToGridCoord(in DDGIUniforms ddgi, int index)
{
    ivec3 iPos;

    // Slow, but works for any # of probes
    iPos.x = index % ddgi.probeCounts.x;
    iPos.y = (index % (ddgi.probeCounts.x * ddgi.probeCounts.y)) / ddgi.probeCounts.x;
    iPos.z = index / (ddgi.probeCounts.x * ddgi.probeCounts.y);

    // Assumes probeCounts are powers of two.
    // Saves ~10ms compared to the divisions above
    // Precomputing the MSB actually slows this code down substantially
    //    iPos.x = index & (ddgi.probeCounts.x - 1);
    //    iPos.y = (index & ((ddgi.probeCounts.x * ddgi.probeCounts.y) - 1)) >> findMSB(ddgi.probeCounts.x);
    //    iPos.z = index >> findMSB(ddgi.probeCounts.x * ddgi.probeCounts.y);

    return iPos;
}



vec3 ProbeLocation(in DDGIUniforms ddgi, int index)
{
    // Compute grid coord from instance ID.
    ivec3 gridCoord = ProbeIndexToGridCoord(ddgi, index);

    // Compute probe position from grid coord.
    return GridCoordToPosition(ddgi, gridCoord);
}



float Square(float v)
{
    return v * v;
}



vec3 Square(vec3 v)
{
    return v * v;
}



float Pow3(float x) 
{ 
    return x * x * x; 
}



float SignNotZero(in float k)
{
    return (k >= 0.0) ? 1.0 : -1.0;
}



vec2 SignNotZero(in vec2 v)
{
    return vec2(SignNotZero(v.x), SignNotZero(v.y));
}



vec2 OctEncode(in vec3 v) 
{
    float l1norm = abs(v.x) + abs(v.y) + abs(v.z);
    vec2 result = v.xy * (1.0 / l1norm);
    if (v.z < 0.0)
        result = (1.0 - abs(result.yx)) * SignNotZero(result.xy);
    return result;
}



vec3 OctDecode(vec2 o)
{
    vec3 v = vec3(o.x, o.y, 1.0 - abs(o.x) - abs(o.y));

    if (v.z < 0.0)
        v.xy = (1.0 - abs(v.yx)) * SignNotZero(v.xy);

    return normalize(v);
}



int ProbeId(vec2 texelXY, int fullTextureWidth, int probeSideLength)
{
    int probeWithBorderSide = probeSideLength + 2;
    int probesPerSide = (fullTextureWidth - 2) / probeWithBorderSide;
    return int(texelXY.x / probeWithBorderSide) + probesPerSide * int(texelXY.y / probeWithBorderSide);
}



// Compute normalized oct coord, mapping top left of top left pixel to (-1,-1)
vec2 NormalizedOctCoord(ivec2 fragCoord, int probeSideLength)
{
    int probeWithBorderSide = probeSideLength + 2;

    vec2 octFragCoord = ivec2((fragCoord.x - 2) % probeWithBorderSide, (fragCoord.y - 2) % probeWithBorderSide);
    // Add back the half pixel to get pixel center normalized coordinates
    return (vec2(octFragCoord) + vec2(0.5f)) * (2.0f / float(probeSideLength)) - vec2(1.0f, 1.0f);
}



vec2 TextureCoordFromDirection(vec3 dir, int probeIndex, int fullTextureWidth, int fullTextureHeight, int probeSideLength) 
{
    vec2 normalizedOctCoord = OctEncode(normalize(dir));
    vec2 normalizedOctCoordZeroOne = (normalizedOctCoord + vec2(1.0f)) * 0.5f;

    // Length of a probe side, plus one pixel on each edge for the border
    float probeWithBorderSide = float(probeSideLength) + 2.0f;

    vec2 octCoordNormalizedToTextureDimensions = (normalizedOctCoordZeroOne * float(probeSideLength)) / vec2(float(fullTextureWidth), float(fullTextureHeight));

    int probesPerRow = (fullTextureWidth - 2) / int(probeWithBorderSide);

    // Add (2,2) back to texCoord within larger texture. Compensates for 1 pix 
    // border around texture and further 1 pix border around top left probe.
    vec2 probeTopLeftPosition = vec2(mod(float(probeIndex), float(probesPerRow)) * probeWithBorderSide,
        (probeIndex / probesPerRow) * probeWithBorderSide) + vec2(2.0f, 2.0f);

    vec2 normalizedProbeTopLeftPosition = vec2(probeTopLeftPosition) / vec2(float(fullTextureWidth), float(fullTextureHeight));

    return vec2(normalizedProbeTopLeftPosition + octCoordNormalizedToTextureDimensions);
}



vec3 SampleIrradiance(in DDGIUniforms ddgi, vec3 P, vec3 N, vec3 Wo, sampler2D irradianceTexture, sampler2D depthTexture)
{
    ivec3 baseGridCoord = BaseGridCoord(ddgi, P);
    vec3 baseProbePos = GridCoordToPosition(ddgi, baseGridCoord);
    
    vec3  sumIrradiance = vec3(0.0f);
    float sumWeight = 0.0f;

    // alpha is how far from the floor(currentVertex) position. on [0, 1] for each axis.
    vec3 alpha = clamp((P - baseProbePos) / ddgi.gridStep, vec3(0.0f), vec3(1.0f));

    // Iterate over adjacent probe cage
    for (int i = 0; i < 8; ++i) 
    {
        // Compute the offset grid coord and clamp to the probe grid boundary
        // Offset = 0 or 1 along each axis
        ivec3 offset = ivec3(i, i >> 1, i >> 2) & ivec3(1);
        ivec3 probeGridCoord = clamp(baseGridCoord + offset, ivec3(0), ddgi.probeCounts - ivec3(1));
        int p = GridCoordToProbeIndex(ddgi, probeGridCoord);

        // Make cosine falloff in tangent plane with respect to the angle from the surface to the probe so that we never
        // test a probe that is *behind* the surface.
        // It doesn't have to be cosine, but that is efficient to compute and we must clip to the tangent plane.
        vec3 probePos = GridCoordToPosition(ddgi, probeGridCoord);

        // Bias the position at which visibility is computed; this
        // avoids performing a shadow test *at* a surface, which is a
        // dangerous location because that is exactly the line between
        // shadowed and unshadowed. If the normal bias is too small,
        // there will be light and dark leaks. If it is too large,
        // then samples can pass through thin occluders to the other
        // side (this can only happen if there are MULTIPLE occluders
        // near each other, a wall surface won't pass through itself.)
        vec3 probeToPoint = P - probePos + (N + 3.0 * Wo) * ddgi.normalBias;
        vec3 dir = normalize(-probeToPoint);

        // Compute the trilinear weights based on the grid cell vertex to smoothly
        // transition between probes. Avoid ever going entirely to zero because that
        // will cause problems at the border probes. This isn't really a lerp. 
        // We're using 1-a when offset = 0 and a when offset = 1.
        vec3 trilinear = mix(1.0 - alpha, alpha, offset);
        float weight = 1.0;

        // Clamp all of the multiplies. We can't let the weight go to zero because then it would be 
        // possible for *all* weights to be equally low and get normalized
        // up to 1/n. We want to distinguish between weights that are 
        // low because of different factors.

        // Smooth backface test
        {
            // Computed without the biasing applied to the "dir" variable. 
            // This test can cause reflection-map looking errors in the image
            // (stuff looks shiny) if the transition is poor.
            vec3 trueDirectionToProbe = normalize(probePos - P);

            // The naive soft backface weight would ignore a probe when
            // it is behind the surface. That's good for walls. But for small details inside of a
            // room, the normals on the details might rule out all of the probes that have mutual
            // visibility to the point. So, we instead use a "wrap shading" test below inspired by
            // NPR work.
            // weight *= max(0.0001, dot(trueDirectionToProbe, wsN));

            // The small offset at the end reduces the "going to zero" impact
            // where this is really close to exactly opposite
            weight *= Square(max(0.0001, (dot(trueDirectionToProbe, N) + 1.0) * 0.5)) + 0.2;
        }

        // Moment visibility test
        if (ddgi.visibilityTest == 1)
        {
            vec2 texCoord = TextureCoordFromDirection(-dir, p, ddgi.depthTextureWidth, ddgi.depthTextureHeight, ddgi.depthProbeSideLength);

            float distToProbe = length(probeToPoint);

            vec2 temp = textureLod(depthTexture, texCoord, 0.0f).rg;
            float mean = temp.x;
            float variance = abs(Square(temp.x) - temp.y);

            // http://www.punkuser.net/vsm/vsm_paper.pdf; equation 5
            // Need the max in the denominator because biasing can cause a negative displacement
            float chebyshevWeight = variance / (variance + Square(max(distToProbe - mean, 0.0)));
                
            // Increase contrast in the weight 
            chebyshevWeight = max(Pow3(chebyshevWeight), 0.0);

            weight *= (distToProbe <= mean) ? 1.0 : chebyshevWeight;
        }

        // Avoid zero weight
        weight = max(0.000001, weight);
                 
        vec3 irradianceDir = N;

        vec2 texCoord = TextureCoordFromDirection(normalize(irradianceDir), p, ddgi.irradianceTextureWidth, ddgi.irradianceTextureHeight, ddgi.irradianceProbeSideLength);

        vec3 probeIrradiance = textureLod(irradianceTexture, texCoord, 0.0f).rgb;

        // A tiny bit of light is really visible due to log perception, so
        // crush tiny weights but keep the curve continuous. This must be done
        // before the trilinear weights, because those should be preserved.
        const float crushThreshold = 0.2f;
        if (weight < crushThreshold)
            weight *= weight * weight * (1.0f / Square(crushThreshold)); 

        // Trilinear weights
        weight *= trilinear.x * trilinear.y * trilinear.z;

        // Weight in a more-perceptual brightness space instead of radiance space.
        // This softens the transitions between probes with respect to translation.
        // It makes little difference most of the time, but when there are radical transitions
        // between probes this helps soften the ramp.
#       if LINEAR_BLENDING == 0
            probeIrradiance = sqrt(probeIrradiance);
#       endif
        
        sumIrradiance += weight * probeIrradiance;
        sumWeight += weight;
    }

    vec3 netIrradiance = sumIrradiance / sumWeight;
    
    netIrradiance.x = isnan(netIrradiance.x) ? 0.5f : netIrradiance.x;
    netIrradiance.y = isnan(netIrradiance.y) ? 0.5f : netIrradiance.y;
    netIrradiance.z = isnan(netIrradiance.z) ? 0.5f : netIrradiance.z;

    // Go back to linear irradiance
#   if LINEAR_BLENDING == 0
        netIrradiance = Square(netIrradiance);
#   endif
    netIrradiance *= ddgi.energyPreservation;

    return 0.5f * M_PI * netIrradiance;
}



#endif